from ultralytics import YOLO
import cv2
import time
import os
import torch
from ultralytics.engine.results import Results
from ultralytics.utils import ops

# 初始化模型
model_path = "./merge_seg_rknn_model"  # RKNN 模型文件夹路径
try:
    model = YOLO(model_path, task="segment")
except Exception as e:
    print(f"加载模型失败: {e}")
    print(f"检查模型文件夹: {os.listdir(model_path) if os.path.exists(model_path) else '文件夹不存在'}")
    exit(1)

# 加载图像
img_path = "./data/Color/st.jpg"
orig_img = cv2.imread(img_path)
if orig_img is None:
    print(f"无法加载图像: {img_path}")
    exit(1)

print(f"原始图像尺寸: {orig_img.shape}")

# 预测（包含前处理、推理、后处理）
try:
    print("\n开始预测...")
    start_total = time.time()
    
    # 使用 model.predict 进行完整预测
    results = model.predict(
        source=orig_img,
        imgsz=640,
        conf=0.25,
        iou=0.45,
        device="cpu",  # RKNN 使用 CPU 模拟
        verbose=False
    )
    
    # 分阶段耗时（手动计时）
    print("\n开始前处理...")
    start_pre = time.time()
    predictor = model.predictor
    im = [orig_img]
    preprocessed = predictor.preprocess(im)
    end_pre = time.time()
    print(f"前处理结果: {preprocessed.shape}")
    print(f"前处理耗时: {(end_pre - start_pre)*1000:.2f}ms")
    
    print("\n开始推理...")
    start_inf = time.time()
    with torch.no_grad():
        pred = predictor.model(preprocessed)
    end_inf = time.time()
    print(f"推理结果: {[p.shape if isinstance(p, torch.Tensor) else type(p) for p in pred] if isinstance(pred, (list, tuple)) else pred.shape}")
    print(f"推理耗时: {(end_inf - start_inf)*1000:.2f}ms")
    
    print("\n开始后处理...")
    start_post = time.time()
    
    # 检查 results 是否有效
    if results and hasattr(results[0], 'masks') and results[0].masks is not None:
        results[0].save(filename="output_image.jpg")
        print("后处理结果已保存为 output_image.jpg")
    else:
        print("警告: model.predict 未生成有效掩码，尝试手动后处理...")
        # 手动后处理 [det, proto]
        if isinstance(pred, (list, tuple)) and len(pred) == 2:
            det, proto = pred  # det: [1, 37, 8400], proto: [1, 32, 160, 160]
            det = det.transpose(-2, -1)  # [1, 37, 8400] -> [1, 8400, 37]
            boxes = det[..., :4]  # [1, 8400, 4]
            scores = det[..., 4]  # [1, 8400]
            classes = det[..., 5:6].argmax(-1)  # [1, 8400], nc=1
            masks_in = det[..., 6:38]  # [1, 8400, 32]
            
            # 非极大值抑制 (NMS)
            indices = ops.non_max_suppression(
                det,
                conf_thres=0.25,
                iou_thres=0.45,
                classes=None,
                agnostic=False,
                max_det=100,
                nc=1  # 类别数为 1
            )
            
            # 处理分割掩码
            if len(indices[0]) > 0:
                boxes = boxes[0, indices[0]]  # [N, 4]
                scores = scores[0, indices[0]]  # [N]
                classes = classes[0, indices[0]]  # [N]
                masks_in = masks_in[0, indices[0]]  # [N, 32]
                proto = proto[0]  # [32, 160, 160]
                
                # 生成掩码
                masks = ops.process_mask(
                    proto=proto,
                    masks_in=masks_in,
                    bboxes=boxes,
                    shape=(640, 640),
                    upsample=True
                )
                
                # 缩放到原始图像尺寸
                orig_shape = orig_img.shape[:2]  # (720, 1280)
                masks = ops.scale_masks(masks, orig_shape, (640, 640))
                boxes = ops.scale_boxes(boxes, orig_shape, (640, 640))
                
                # 构造 Results 对象
                results = [Results(
                    orig_img=orig_img,
                    boxes=boxes,
                    masks=masks,
                    confs=scores,
                    cls=classes,
                    names={0: 'battery-top'}
                )]
                results[0].save(filename="output_image.jpg")
                print("手动后处理结果已保存为 output_image.jpg")
            else:
                cv2.imwrite("output_image.jpg", orig_img)
                print("未检测到有效目标，保存原始图像...")
        else:
            cv2.imwrite("output_image.jpg", orig_img)
            print("推理输出格式不正确，保存原始图像...")
    
    end_post = time.time()
    print(f"后处理耗时: {(end_post - start_post)*1000:.2f}ms")
    print(f"总耗时: {(end_post - start_total)*1000:.2f}ms")
except Exception as e:
    print(f"预测过程中出错: {e}")
    import traceback
    traceback.print_exc()
    exit(1)